/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.dubbo.rpc.cluster.loadbalance;

import org.apache.dubbo.common.URL;
import org.apache.dubbo.common.io.Bytes;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.Invoker;
import org.apache.dubbo.rpc.support.RpcUtils;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import static org.apache.dubbo.common.constants.CommonConstants.COMMA_SPLIT_PATTERN;

/**
 * ConsistentHashLoadBalance
 */
public class ConsistentHashLoadBalance extends AbstractLoadBalance {
    public static final String NAME = "consistenthash";

    /**
     * Hash nodes name
     */
    public static final String HASH_NODES = "hash.nodes";

    /**
     * Hash arguments name
     */
    public static final String HASH_ARGUMENTS = "hash.arguments";

    private final ConcurrentMap<String, ConsistentHashSelector<?>> selectors = new ConcurrentHashMap<String, ConsistentHashSelector<?>>();

    /**
     * 一致哈希负载平衡
     *      1. 一致性 Hash，相同参数的请求总是发到同一提供者。
     *      2. 当某一台提供者挂时，原本发往该提供者的请求，基于虚拟节点，平摊到其它提供者，不会引起剧烈变动
     *      3. 算法参见：Consistent Hashing | WIKIPEDIA
     *      4. 缺省只对第一个参数 Hash，如果要修改，请配置 <dubbo:parameter key="hash.arguments" value="0,1" />
     *      5. 缺省用 160 份虚拟节点，如果要修改，请配置 <dubbo:parameter key="hash.nodes" value="320" />
     * @param invokers
     * @param url
     * @param invocation
     * @return
     * @param <T>
     */
    @SuppressWarnings("unchecked")
    @Override
    protected <T> Invoker<T> doSelect(List<Invoker<T>> invokers, URL url, Invocation invocation) {
        // 获取方法名称
        String methodName = RpcUtils.getMethodName(invocation);

        String key = invokers.get(0).getUrl().getServiceKey() + "." + methodName;

        // using the hashcode of list to compute the hash only pay attention to the elements in the list
        // 获取 invoker列表哈希值
        int invokersHashCode = invokers.hashCode();

        // 获取 一致性哈希选择器
        ConsistentHashSelector<T> selector = (ConsistentHashSelector<T>) selectors.get(key);

        if (selector == null || selector.identityHashCode != invokersHashCode) {
            // 创建一致性哈希选择器（核心）
            selectors.put(key, new ConsistentHashSelector<T>(invokers, methodName, invokersHashCode));
            selector = (ConsistentHashSelector<T>) selectors.get(key);
        }
        return selector.select(invocation);
    }

    private static final class ConsistentHashSelector<T> {

        private final TreeMap<Long, Invoker<T>> virtualInvokers;

        private final int replicaNumber;

        private final int identityHashCode;

        private final int[] argumentIndex;

        ConsistentHashSelector(List<Invoker<T>> invokers, String methodName, int identityHashCode) {

            this.virtualInvokers = new TreeMap<Long, Invoker<T>>();

            this.identityHashCode = identityHashCode;

            URL url = invokers.get(0).getUrl();

            // 获取 hash.nodes 参数值，默认 160
            this.replicaNumber = url.getMethodParameter(methodName, HASH_NODES, 160);
            // 获取 hash.arguments 参数值。默认 0
            String[] index = COMMA_SPLIT_PATTERN.split(url.getMethodParameter(methodName, HASH_ARGUMENTS, "0"));

            // 参数索引
            argumentIndex = new int[index.length];

            for (int i = 0; i < index.length; i++) {

                argumentIndex[i] = Integer.parseInt(index[i]);
            }

            //  初始化hash环，将所有的invoker均匀的分布到环上
            for (Invoker<T> invoker : invokers) {
                // 获取地址
                String address = invoker.getUrl().getAddress();
                //  创建虚拟节点，为每个物理机节点应该有replicaNumber个虚拟节点，但是这里用两次for循环，应该是为了加强hash的随机性
                for (int i = 0; i < replicaNumber / 4; i++) {
                    // 加密
                    byte[] digest = Bytes.getMD5(address + i);

                    for (int h = 0; h < 4; h++) {
                        // 计算hash值
                        long m = hash(digest, h);
                        // 添加 提供者虚拟节点
                        virtualInvokers.put(m, invoker);
                    }
                }
            }
        }

        public Invoker<T> select(Invocation invocation) {
            String methodName = RpcUtils.getMethodName(invocation);
            byte[] digest = Bytes.getMD5(methodName);
            long hash = hash(digest, 0);
            return selectForKey(hash);
        }

        private String toKey(Object[] args, boolean isGeneric) {
            return isGeneric ? toKey((Object[]) args[1]) : toKey(args);
        }

        private String toKey(Object[] args) {
            StringBuilder buf = new StringBuilder();
            for (int i : argumentIndex) {
                if (i >= 0 && args != null && i < args.length) {
                    buf.append(args[i]);
                }
            }
            return buf.toString();
        }

        private Invoker<T> selectForKey(long hash) {

            Map.Entry<Long, Invoker<T>> entry = virtualInvokers.ceilingEntry(hash);

            if (entry == null) {
                // 获取第一个
                entry = virtualInvokers.firstEntry();
            }

            return entry.getValue();
        }

        private long hash(byte[] digest, int number) {
            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;
        }
    }

}
